package cn.gwm.flink.streaming.source;//
// Source code recreated from a .class file by IntelliJ IDEA
// (powered by FernFlower decompiler)
//

import com.alibaba.ververica.connectors.common.source.AbstractDynamicParallelSource;
import com.alibaba.ververica.connectors.common.source.SourceUtils;
import com.alibaba.ververica.connectors.common.source.reader.RecordReader;
import com.alibaba.ververica.connectors.datahub.DatahubClientProvider;
import com.alibaba.ververica.connectors.datahub.source.DatahubRecordReader;
import com.alibaba.ververica.connectors.datahub.source.DatahubShardInputSplit;
import com.aliyun.datahub.client.DatahubClient;
import com.aliyun.datahub.client.http.HttpConfig;
import com.aliyun.datahub.client.model.ListShardResult;
import com.aliyun.datahub.client.model.RecordEntry;
import com.aliyun.datahub.client.model.ShardEntry;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Comparator;
import java.util.Iterator;
import java.util.LinkedList;
import java.util.List;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.core.io.InputSplit;
import org.apache.flink.util.StringUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * 自定义datahub 读取数据类
 * @ClassName DatahubRecordsReader
 * @Author xzh
 * @Date 2023/03/09 11:10
 **/
public class DatahubSourcesFunction extends AbstractDynamicParallelSource<RecordEntry, Long> {
    private static final long serialVersionUID = 1L;
    private static Logger logger = LoggerFactory.getLogger(DatahubSourcesFunction.class);
    private String endpoint;
    private String projectName;
    private String topicName;
    private String subId;
    private String accessId;
    private String accessKey;
    private Configuration properties;
    private long startTimeInMs;
    private long stopTimeInMs;
    private int requestTimeout;
    private int maxFetchSize;
    private int maxBufferSize;
    private int retryInterval;
    private int retryTimeout;
    private int fetchLatestDelay;
    private transient DatahubClientProvider clientProvider;
    private List<String> initShardsList;

    public DatahubSourcesFunction(String endpoint, String projectName, String topicName, String subId, String accessId, String accessKey, long stopTimeInMs) {
        this(endpoint, projectName, topicName, subId, accessId, accessKey, new Configuration(), 0L, stopTimeInMs);
    }

    public DatahubSourcesFunction(String endpoint, String projectName, String topicName, String subId, String accessId, String accessKey, long startTimeInMs, long stopTimeInMs) {
        this(endpoint, projectName, topicName, subId, accessId, accessKey, new Configuration(), startTimeInMs, stopTimeInMs);
    }

    public DatahubSourcesFunction(String endpoint, String projectName, String topicName, String subId, String accessId, String accessKey, Configuration properties, long startTimeInMs, long stopTimeInMs) {
        this.accessId = null;
        this.accessKey = null;
        this.requestTimeout = 30000;
        this.maxFetchSize = 2000;
        this.maxBufferSize = 2000;
        this.retryInterval = 1000;
        this.retryTimeout = 1800000;
        this.fetchLatestDelay = 500;
        this.endpoint = endpoint;
        this.projectName = projectName;
        this.topicName = topicName;
        this.subId = subId;
        this.accessId = accessId;
        this.accessKey = accessKey;
        this.startTimeInMs = startTimeInMs;
        this.stopTimeInMs = stopTimeInMs;
        this.properties = properties;
        this.clientProvider = this.createProvider();
    }

    public void setRequestTimeout(int requestTimeout) {
        this.requestTimeout = requestTimeout;
    }

    public void setMaxFetchSize(int maxFetchSize) {
        this.maxFetchSize = maxFetchSize;
    }

    public void setMaxBufferSize(int maxBufferSize) {
        this.maxBufferSize = maxBufferSize;
    }

    public void setRetryInterval(int retryInterval) {
        this.retryInterval = retryInterval;
    }

    public void setRetryTimeout(int retryTimeout) {
        this.retryTimeout = retryTimeout;
    }

    public void setFetchLatestDelay(int fetchLatestDelay) {
        this.fetchLatestDelay = fetchLatestDelay;
    }

    /**
     * 根据是否传入开始消费时间位点
     * @auth xzh
     * @param config
     * @return
     * @throws IOException
     */
    public RecordReader<RecordEntry, Long> createReader(Configuration config) throws IOException {
        this.initShardsList();
        if(this.startTimeInMs==0){
            DatahubRecordsReader reader = !StringUtils.isNullOrWhitespaceOnly(this.accessId) && !StringUtils.isNullOrWhitespaceOnly(this.accessKey) ? new DatahubRecordsReader(this.endpoint, this.projectName, this.topicName, this.subId, this.accessId, this.accessKey, this.initShardsList.size(), this.startTimeInMs, this.stopTimeInMs,this.initShardsList) : new DatahubRecordsReader(this.endpoint, this.projectName, this.topicName, this.subId, this.properties, this.initShardsList.size(), this.startTimeInMs, this.stopTimeInMs,this.initShardsList);
            reader.setRequestTimeout(this.requestTimeout).setMaxFetchSize(this.maxFetchSize).setMaxBufferSize(this.maxBufferSize).setRetryInterval(this.retryInterval).setRetryTimeout(this.retryTimeout).setFetchLatestDelay(this.fetchLatestDelay).setClientProvider(this.clientProvider);
            return reader;
        }else{
            DatahubRecordReader reader = !StringUtils.isNullOrWhitespaceOnly(this.accessId) && !StringUtils.isNullOrWhitespaceOnly(this.accessKey) ? new DatahubRecordReader(this.endpoint, this.projectName, this.topicName, this.subId, this.accessId, this.accessKey, this.initShardsList.size(), this.startTimeInMs, this.stopTimeInMs) : new DatahubRecordReader(this.endpoint, this.projectName, this.topicName, this.subId, this.properties, this.initShardsList.size(), this.startTimeInMs, this.stopTimeInMs);
            reader.setRequestTimeout(this.requestTimeout).setMaxFetchSize(this.maxFetchSize).setMaxBufferSize(this.maxBufferSize).setRetryInterval(this.retryInterval).setRetryTimeout(this.retryTimeout).setFetchLatestDelay(this.fetchLatestDelay).setClientProvider(this.clientProvider);
            return reader;
        }
    }

    public InputSplit[] createInputSplitsForCurrentSubTask(int numberOfParallelSubTasks, int indexOfThisSubTask) throws IOException {
        this.initShardsList();
        int totalPartitionCount = this.initShardsList.size();
        List<Integer> subscribedShardIndexList = SourceUtils.modAssign("datahub-" + this.projectName + ":" + this.topicName, numberOfParallelSubTasks, indexOfThisSubTask, totalPartitionCount);
        DatahubShardInputSplit[] inputSplits = new DatahubShardInputSplit[subscribedShardIndexList.size()];
        int i = 0;

        Integer shardIndex;
        for(Iterator var7 = subscribedShardIndexList.iterator(); var7.hasNext(); inputSplits[i++] = new DatahubShardInputSplit(shardIndex, (String)this.initShardsList.get(shardIndex))) {
            shardIndex = (Integer)var7.next();
        }

        return inputSplits;
    }

    public List<Tuple2<InputSplit, Long>> reAssignInputSplitsForCurrentSubTask(int numberOfParallelSubTasks, int indexOfThisSubTask, List<AbstractDynamicParallelSource.InnerProgress<Long>> allSplitsInState) throws IOException {
        this.initShardsList();
        List<Tuple2<InputSplit, Long>> initProgress = new ArrayList();
        List<String> subscribedPartitions = this.modAssign(numberOfParallelSubTasks, indexOfThisSubTask);
        Iterator var6 = subscribedPartitions.iterator();

        while(var6.hasNext()) {
            String shard = (String)var6.next();
            boolean existBefore = false;
            Iterator var9 = allSplitsInState.iterator();

            while(var9.hasNext()) {
                AbstractDynamicParallelSource.InnerProgress<Long> progress = (AbstractDynamicParallelSource.InnerProgress)var9.next();
                DatahubShardInputSplit split = (DatahubShardInputSplit)progress.getInputSplit();
                if (Integer.parseInt(shard) == Integer.parseInt(split.getShardId())) {
                    logger.info("Current subTask [{}], split [{}], progress: [{}]", new Object[]{indexOfThisSubTask, shard, progress.getCursor()});
                    initProgress.add(new Tuple2(progress.getInputSplit(), progress.getCursor()));
                    existBefore = true;
                    break;
                }
            }

            if (!existBefore) {
                DatahubShardInputSplit inputSplit = new DatahubShardInputSplit(Integer.parseInt(shard), shard);
                logger.info("Current subTask [{}], new split [{}]", indexOfThisSubTask, inputSplit);
                initProgress.add(Tuple2.of(inputSplit, -1L));
            }
        }

        return initProgress;
    }

    public List<String> getPartitionList() throws Exception {
        List<String> shardIds = new ArrayList();
        if (null == this.clientProvider) {
            this.clientProvider = this.createProvider();
        }

        ListShardResult shardResult = ((DatahubClient)this.clientProvider.getClient()).listShard(this.projectName, this.topicName);
        Iterator var3 = shardResult.getShards().iterator();

        while(var3.hasNext()) {
            ShardEntry shardEntry = (ShardEntry)var3.next();
            shardIds.add(shardEntry.getShardId());
        }

        shardIds.sort(Comparator.comparingInt(Integer::parseInt));
        return shardIds;
    }

    private void initShardsList() {
        if (null == this.initShardsList || this.initShardsList.size() == 0) {
            try {
                this.initShardsList = this.getPartitionList();
            } catch (Exception var2) {
                logger.error("init shard list failed", var2);
                throw new RuntimeException(var2);
            }
        }

    }

    public String toString() {
        return String.format("%s:%s:%s", this.getClass().getSimpleName(), this.projectName, this.topicName);
    }

    private DatahubClientProvider createProvider() {
        HttpConfig httpConfig = (new HttpConfig()).setConnTimeout(this.requestTimeout).setReadTimeout(this.requestTimeout);
        return new DatahubClientProvider(this.endpoint, this.accessId, this.accessKey, this.properties, httpConfig);
    }

    private List<String> modAssign(int consumerCount, int consumerIndex) {
        List<String> assignedShards = new LinkedList();
        Iterator var4 = this.initShardsList.iterator();

        while(var4.hasNext()) {
            String shard = (String)var4.next();
            if (Integer.parseInt(shard) % consumerCount == consumerIndex) {
                assignedShards.add(shard);
            }
        }

        return assignedShards;
    }
}
